package com.jingge.home.framework;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jingge.home.dto.ConversationDto;
import com.jingge.home.util.LocalCacheManager;
import com.jingge.home.util.SpringUtil;
import com.jingge.home.util.SyncHelper;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;

import java.net.URLDecoder;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

/**
 * Dify流式对话的客户端
 *
 * @author Yongxiang
 * @date 2025-06-04 17:25
 * @email 1158055613@qq.com
 */
@Slf4j
@SuppressWarnings("unchecked")
public class DifyConversationStreamChatClient {
	private static final OkHttpClient client =
		new OkHttpClient.Builder().connectTimeout(3, TimeUnit.MINUTES).readTimeout(5, TimeUnit.MINUTES)
		                          .writeTimeout(5, TimeUnit.MINUTES).build();
	private static final ObjectMapper objectMapper = new ObjectMapper();

	/**
	 * 发送对话消息到Dify
	 *
	 * @param query                 消息内容
	 * @param aiUrl                 Dify的URL
	 * @param picUrl                如果有图片，当前图片的URL
	 * @param currentConversationId 当前对话ID
	 * @param user                  用户
	 * @param apiKey                Dify工作流的API KEY
	 * @param timeout               超时时间，超时后直接返回已生成的消息
	 * @param receiveMsgId          微信消息ID
	 * @param eventHandler          流式消息事件的函数接口处理器
	 * @author Yongxiang(* ^ ▽ ^ *)
	 * @date 2025/7/30 14:34
	 */
	public static void sendStreamingMessage(String query,
	                                        String aiUrl,
	                                        String picUrl,
	                                        String currentConversationId,
	                                        String user,
	                                        String apiKey,
	                                        long timeout,
	                                        String receiveMsgId,
	                                        Consumer<StreamEvent> eventHandler) {

		startHandingMsg(user);
		ConversationDto dto = StringUtils.isEmpty(picUrl) ? new ConversationDto(query, currentConversationId, user) :
			new ConversationDto(query, currentConversationId, user, picUrl);
		String requestBody = dto.toString();

		Request request =
			new Request.Builder().url(aiUrl).post(RequestBody.create(requestBody, MediaType.get("application/json")))
			                     .addHeader("Authorization", "Bearer " + apiKey)
			                     .addHeader("Content-Type", "application/json").build();
		log.info("请求：{}，请求体：{}", request, requestBody);
		Response response = null;
		try {
			response = client.newCall(request).execute();
			log.info("开始处理响应：{}", response);
			if (!response.isSuccessful()) {
				eventHandler.accept(new StreamEvent("error", null, "HTTP " + response.code()));
				return;
			}
			ResponseBody body = response.body();
			if (body == null) {
				return;
			}
			LocalCacheManager<String, String> msgCache = SpringUtil.getNonnullBean(LocalCacheManager.class);
			StringBuilder allMsg = new StringBuilder();
			ThreadPoolExecutor msgHandleThreadPool = SpringUtil.getNonnullBean(ThreadPoolExecutor.class);
			Response finalResponse = response;
			AtomicBoolean finished = new AtomicBoolean(false);
			Callable<Boolean> runnable = () -> {
				try {
					log.info("线程开始处理AI流式消息响应...");
					//消息开始处理
					msgCache.set(CacheKeyConstant.MSG_ANSWER_FINISHED_PREFIX + receiveMsgId, "0", 5, TimeUnit.MINUTES);
					String l = "";
					while (!finished.get()) {
						l = body.source().readUtf8Line();
						if (l == null) {
							continue;
						}
						if (l.startsWith("data:")) {
							String eventData = l.substring(5).trim();
							try {
								JsonNode node = objectMapper.readTree(eventData);
								String eventType = node.path("event").asText();
								String taskId = node.path("task_id").asText(null);
								String messageId = node.path("message_id").asText(null);
								String conversationId = node.path("conversation_id").asText(null);

								// 根据不同事件类型处理
								switch (eventType) {
									case "workflow_started":
									case "node_started":
									case "node_finished":
										if (eventType.equals("workflow_started")) {
											log.info("工作流开始");
										}
										eventHandler.accept(new StreamEvent(eventType, taskId, eventData));
										break;
									case "workflow_finished":
										log.info("工作流结束，退出流式消息读取");
										finished.set(true);
										break;
									case "message":
										String raw = node.path("answer").asText();
										log.info("原始消息：{}", raw);
										String answer = "";
										try {
											answer = URLDecoder.decode(raw, "utf-8");
										} catch (Exception e) {
											log.info("raw解码异常", e);
											answer = raw;
										}
										log.info("解析到消息：{}", answer);
										allMsg.append(answer);
										msgCache.set(receiveMsgId, allMsg.toString(), 5, TimeUnit.MINUTES);
										msgCache.set(CacheKeyConstant.RECENT_MSG_ID_PREFIX + user, receiveMsgId, 5, TimeUnit.MINUTES);
										eventHandler.accept(new MessageEvent(taskId, messageId, conversationId, answer));
										break;
									case "message_end":
										JsonNode metadata = node.path("metadata");
										eventHandler.accept(
											new MessageEndEvent(taskId, messageId, conversationId, metadata));
										log.info("消息结束");
										break;
									case "tts_message":
										String audio = node.path("audio").asText();
										eventHandler.accept(
											new TtsMessageEvent(taskId, messageId, conversationId, audio));
										break;
									case "error":
										String errorMsg = node.path("message").asText();
										eventHandler.accept(new StreamEvent("error", taskId, errorMsg));
										log.info("消息错误，退出流式消息读取");
										finished.set(true);
										break;
									default:
										eventHandler.accept(new StreamEvent(eventType, taskId, eventData));
								}
							} catch (Exception e) {
								eventHandler.accept(
									new StreamEvent("parse_error", null, "Failed to parse: " + eventData));
								break;
							}
						}
					}
				} catch (Exception e) {
					log.info("解析消息时发生异常", e);
					eventHandler.accept(new StreamEvent("parse_error", null, "Failed to parse: " + e.getMessage()));
				} finally {
					IOUtils.closeQuietly(body);
					IOUtils.closeQuietly(finalResponse);
					endHandingMsg(user);
				}
				log.info("线程已处理完响应");
				msgCache.set(CacheKeyConstant.MSG_ANSWER_FINISHED_PREFIX + receiveMsgId, "1", 5, TimeUnit.MINUTES);
				return true;
			};
			try {
				Future<Boolean> future = msgHandleThreadPool.submit(runnable);
				SyncHelper.aWaitQuietly(future::isDone, timeout * 1000 - 1500, TimeUnit.MILLISECONDS);
				if (!future.isDone()) {
					eventHandler.accept(new StreamEvent("timeoutMsg", null, receiveMsgId));
				}
			} catch (RejectedExecutionException e) {
				log.info("线程池已满载");
				eventHandler.accept(new StreamEvent("override", null, null));
			}
		} catch (Exception e) {
			log.info("处理响应发生异常", e);
		}

	}

	/**
	 * 判断某个MsgID的消息是否已经处理完成
	 *
	 * @param receiveMsgId 消息ID
	 * @return boolean
	 * @author Yongxiang(* ^ ▽ ^ *)
	 * @date 2025/7/30 10:57
	 */
	public static boolean isMsgIdFinished(String receiveMsgId){
		return "1".equals(SpringUtil.getNonnullBean(LocalCacheManager.class)
				.getAsString(CacheKeyConstant.MSG_ANSWER_FINISHED_PREFIX + receiveMsgId));
	}

	public static boolean isUserMsgHandling(String user) {
		return SpringUtil.getNonnullBean(LocalCacheManager.class).exist(CacheKeyConstant.CONVERSATION_IN_HANDLING_PREFIX + user);
	}

	/**
	 * 标记开始处理用户的消息
	 *
	 * @param user 用户
	 * @author Yongxiang(* ^ ▽ ^ *)
	 * @date 2025/7/30 14:33
	 */
	private static void startHandingMsg(String user) {
		SpringUtil.getNonnullBean(LocalCacheManager.class).set(CacheKeyConstant.CONVERSATION_IN_HANDLING_PREFIX + user, "1", 3, TimeUnit.MINUTES);
	}

	/**
	 * 某个用户消息处理完成
	 * @param user    用户
	 * @author Yongxiang(* ^ ▽ ^ *)
	 * @date 2025/7/30 14:33
	 */
	private static void endHandingMsg(String user) {
		SpringUtil.getNonnullBean(LocalCacheManager.class).remove(CacheKeyConstant.CONVERSATION_IN_HANDLING_PREFIX + user);
	}

	// 基础事件类
	@Getter
	public static class StreamEvent {
		private final String eventType;
		private final String taskId;
		private final String rawData;

		public StreamEvent(String eventType,
		                   String taskId,
		                   String rawData) {
			this.eventType = eventType;
			this.taskId = taskId;
			this.rawData = rawData;
		}

	}

	// 消息事件（特殊处理）
	@Getter
	public static class MessageEvent extends StreamEvent {
		private final String messageId;
		private final String conversationId;
		private final String answer;

		public MessageEvent(String taskId,
		                    String messageId,
		                    String conversationId,
		                    String answer) {
			super("message", taskId, null);
			this.messageId = messageId;
			this.conversationId = conversationId;
			this.answer = answer;
		}

	}

	// 消息结束事件
	@Getter
	public static class MessageEndEvent extends StreamEvent {
		private final String messageId;
		private final String conversationId;
		private final JsonNode metadata;

		public MessageEndEvent(String taskId,
		                       String messageId,
		                       String conversationId,
		                       JsonNode metadata) {
			super("message_end", taskId, null);
			this.messageId = messageId;
			this.conversationId = conversationId;
			this.metadata = metadata;
		}

	}

	// TTS消息事件
	@Getter
	public static class TtsMessageEvent extends StreamEvent {
		private final String messageId;
		private final String conversationId;
		private final String audio;

		public TtsMessageEvent(String taskId,
		                       String messageId,
		                       String conversationId,
		                       String audio) {
			super("tts_message", taskId, null);
			this.messageId = messageId;
			this.conversationId = conversationId;
			this.audio = audio;
		}
	}

}